#!/usr/bin/env python3
"""
测试Transformer生成器的基本功能
"""

import torch
import numpy as np
import sys
import os

# 添加models目录到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'models'))

from models.image_generator import ImageGenerator
from models.audio_generator import AudioGenerator
from models.video_generator import VideoGenerator

def test_image_generator():
    """测试图像生成器"""
    print("🧪 测试图像生成器...")
    
    # 创建图像生成器
    image_gen = ImageGenerator(
        image_size=(64, 64),
        patch_size=8,
        embedding_dim=128,  # 使用较小的维度以加快测试
        nhead=4,
        num_layers=2
    )
    
    # 测试从numpy数组生成
    test_image = np.random.rand(64, 64, 3)  # 随机RGB图像
    generated = image_gen.generate(test_image, num_steps=3)
    
    print("✅ 图像生成器测试通过")
    print(f"   输入形状: {test_image.shape}")
    print(f"   输出形状: {generated.shape}")
    
    # 使用assert而不是return
    assert generated is not None
    assert generated.shape == test_image.shape

def test_audio_generator():
    """测试音频生成器"""
    print("\n🎵 测试音频生成器...")
    
    # 创建音频生成器
    audio_gen = AudioGenerator(
        sample_rate=22050,
        segment_length=512,
        embedding_dim=128,
        nhead=4,
        num_layers=2
    )
    
    # 测试从numpy数组生成
    test_audio = np.random.rand(22050)  # 1秒的随机音频
    generated = audio_gen.generate(test_audio, num_steps=3)
    
    print("✅ 音频生成器测试通过")
    print(f"   输入长度: {len(test_audio)}")
    print(f"   输出形状: {generated.shape}")
    
    # 使用assert而不是return
    assert generated is not None
    assert generated.shape[0] == len(test_audio)

def test_video_generator():
    """测试视频生成器"""
    print("\n🎬 测试视频生成器...")
    
    # 创建视频生成器
    video_gen = VideoGenerator(
        frame_size=(32, 32),
        patch_size=8,
        num_frames=8,
        embedding_dim=128,
        nhead=4,
        num_layers=2
    )
    
    # 测试从numpy数组生成
    test_video = np.random.rand(8, 32, 32, 3)  # 8帧的随机视频
    generated = video_gen.generate(test_video, num_steps=3)
    
    print("✅ 视频生成器测试通过")
    print(f"   输入形状: {test_video.shape}")
    print(f"   输出形状: {generated.shape}")
    
    # 使用assert而不是return
    assert generated is not None
    assert generated.shape == test_video.shape

def main():
    """主测试函数"""
    print("🚀 开始测试Transformer生成器...\n")
    
    # 运行测试
    results = [
        test_image_generator(),
        test_audio_generator(), 
        test_video_generator()
    ]
    
    # 统计结果
    success_count = sum(results)
    total_tests = len(results)
    
    print(f"\n📊 测试完成: {success_count}/{total_tests} 个测试通过")
    
    if success_count == total_tests:
        print("🎉 所有生成器测试通过！")
        print("\n💡 使用建议:")
        print("   1. 运行 'python install_dependencies.py' 安装完整依赖")
        print("   2. 使用真实数据进行训练和生成")
        print("   3. 调整模型参数以获得更好的效果")
    else:
        print("⚠️  部分测试失败，但基本功能仍可使用")
        print("   图像生成器不需要额外依赖即可工作")

if __name__ == "__main__":
    main()